package com.szj.szjaicode.ai;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.szj.szjaicode.ai.guardrail.PromptSafetyInputGuardrail;
import com.szj.szjaicode.ai.guardrail.RetryOutputGuardrail;
import com.szj.szjaicode.ai.tools.*;
import com.szj.szjaicode.exception.BusinessException;
import com.szj.szjaicode.exception.ErrorCode;
import com.szj.szjaicode.model.enums.CodeGenTypeEnum;
import com.szj.szjaicode.service.ChatHistoryService;
import com.szj.szjaicode.utils.SpringContextUtil;
import dev.langchain4j.community.store.memory.chat.redis.RedisChatMemoryStore;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.service.AiServices;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.time.Duration;

/**
 * Ai 服务创建工厂
 */
@Configuration
@Slf4j
public class AiCodeGeneratorServiceFactory {

    @Resource(name = "openAiChatModel")
    private ChatModel chatModel;

    @Resource
    private RedisChatMemoryStore redisChatMemoryStore;

    @Resource
    private ChatHistoryService chatHistoryService;

    @Resource
    private ToolManager toolManager;

    /**
     * AI 服务实例缓存
     * 缓存策略：
     * 1. 最大缓存1000个实例
     * 2. 写入后30分钟过期
     * 3. 访问后10分钟过期
     */
    private final Cache<String, AiCodeGeneratorService> serviceCache = Caffeine.newBuilder()
            .maximumSize(1000)
            .expireAfterWrite(Duration.ofMinutes(30))
            .expireAfterAccess(Duration.ofMinutes(10))
            .removalListener((key, value, cause) -> {
                log.debug("Ai服务实例被移除，缓存键: {}，原因：{}", key, cause);
            })
            .build();

    /**
     * 根据appId获取服务(为了兼容老逻辑)
     *
     * @param appId
     * @return
     * TODO 实现对话记忆
     */
    public AiCodeGeneratorService getAiCodeGeneratorService(long appId) {
       return getAiCodeGeneratorService(appId, CodeGenTypeEnum.HTML);
    }

    /**
     * 根据appId获取服务
     *
     * @param appId
     * @return
     * TODO 实现对话记忆
     */
    public AiCodeGeneratorService getAiCodeGeneratorService(long appId, CodeGenTypeEnum codeGenType) {
        String cacheKey = buildCacheKey(appId, codeGenType);
        return serviceCache.get(cacheKey,key -> createAiCodeGeneratorService(appId, codeGenType));
    }

    /**
     * 创建新的 AI 服务实例
     * @param appId 应用ID
     * @param codeGenType 代码生成类型
     * @return
     */
    private AiCodeGeneratorService createAiCodeGeneratorService(long appId, CodeGenTypeEnum codeGenType) {
        log.info("为appId：{} 创建Ai服务实例", appId);
        //根据appId 构建独立的对话记忆
        //当用户调用该 AI 服务的方法（例如发送消息）时，
        //langchain4j 框架会自动将用户的输入消息添加到chatMemory 中，同时也会将 AI 的响应消息存入记忆。
        MessageWindowChatMemory chatMemory = MessageWindowChatMemory
                .builder()
                .id(appId)
                .chatMemoryStore(redisChatMemoryStore)
                .maxMessages(20)
                .build();
        //从数据库中记载对话历史到记忆中
        chatHistoryService.loadChatHistoryToMemory(appId, chatMemory, 20);

        return switch (codeGenType){
            //HTML和多文件生成，使用流式对话模型
            case HTML,MULTI_FILE -> {
                //使用多例模式的 StreamingChatModel 解决并发问题
                StreamingChatModel openAiStreamingChatModel = SpringContextUtil.getBean("streamingChatModelPrototype", StreamingChatModel.class);
                yield AiServices.builder(AiCodeGeneratorService.class)
                        .chatModel(chatModel)
                        .streamingChatModel(openAiStreamingChatModel)
                        .chatMemory(chatMemory)
                        .inputGuardrails(new PromptSafetyInputGuardrail()) //添加输入护轨
                        //.outputGuardrails(new RetryOutputGuardrail()) //添加输出护轨,为了流式输出，这里不使用
                        .build();
            }
            // Vue项目生成，使用工具调用和推理模型
            case VUE_PROJECT -> {
                //使用多例模式的 StreamingChatModel 解决并发问题
                StreamingChatModel reasoningStreamingChatModel = SpringContextUtil.getBean("reasoningStreamingChatModelPrototype", StreamingChatModel.class);
                yield AiServices.builder(AiCodeGeneratorService.class)
                    .chatModel(chatModel)
                    .streamingChatModel(reasoningStreamingChatModel)
                    .chatMemoryProvider(memoryId -> chatMemory) //这样写避免非空错误（框架的问题）
                    .tools(toolManager.getAllTools())
                    // 处理工具调用幻觉问题
                    .hallucinatedToolNameStrategy(toolExecutionRequest ->
                           ToolExecutionResultMessage.from(toolExecutionRequest,
                                   "Error: there is no tool called" + toolExecutionRequest.name())
                    )
                        .maxSequentialToolsInvocations(20) //最多连续调用工具20次
                        .inputGuardrails(new PromptSafetyInputGuardrail())   //添加输入护轨
                        //.outputGuardrails(new RetryOutputGuardrail()) //添加输出护轨，为了流式输出，这里不使用
                    .build();
            }
            default -> throw new BusinessException(ErrorCode.SYSTEM_ERROR, "不支持的代码生成类型：" + codeGenType.getValue());
        };
    }

    /**
     * 创建 Ai代码生成器服务
     *
     * @return
     */
    @Bean
    public AiCodeGeneratorService aiCodeGeneratorService() {
        return getAiCodeGeneratorService(0);
    }

    /**
     * 构造缓存键
     *
     * @param appId
     * @param codeGenType
     * @return
     */
    private String buildCacheKey(long appId, CodeGenTypeEnum codeGenType){
        return appId + "_" + codeGenType.getValue();
    }
}